#include <sds/sds.h>

#include <cluster.h>
#include <matrix.h>
#include <segment.h>

#include <random.h>
#include <chi_square.h>

#include <vector>
#include <iostream>
#include <iomanip>

using namespace sds;
using namespace std;

Random rng;

// Parameters

double noise = 0.0001;

double leave = 0.01;
double pixelLeave = 0.01;

double min_relevant = 0.25;
double min_rate = 0.25;

double scaleRate = 1.1;

// Constans

double leaveChiSquare = invChiSquare(1.0 - leave, CH);
double pixelLeaveChiSquare = invChiSquare(1.0 - pixelLeave, CH);

enum Todo {
  LEAVE, CLASSIFY, SPLIT, RECLASSIFY
};


int main(int argc, const char* argv[]) {
  SFile segfile(argv[2]);
  SDataView<unsigned char> seg =
    segfile.getDataset(0).getDataView<unsigned char>();

  SFile clufile(argv[3]);
  SDataView<unsigned char> clu =
    clufile.getDataset(0).getDataView<unsigned char>();

  SFile reffile(argv[4]);
  SDataView<unsigned char> ref =
    reffile.getDataset(0).getDataView<unsigned char>();

  vector<SInteger> size = ref.getSize();

  SFile srcfile(argv[1]);
  SDataView<Point> src(size);

  for (int k = 0; k < CH; ++k) {
    SDataView<unsigned char> view =
      srcfile.getDataset(k).getDataView<unsigned char>();
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	((Point&)src[i][j])[k] = (double)view[i][j];
      }    
    }
  }
  

  Cluster signs[256];
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (clu[i][j] != 0) {
	signs[clu[i][j]] += src[i][j];
      }
    }
  }

  vector<unsigned char> index;
  vector<Vector> expected;
  vector<Matrix> inverse;
  vector<double> lndet;

  for (int k = 1; k < 256; ++k) {
    if (signs[k].size() == 0) continue;
    signs[k].scale(scaleRate);
    index.push_back(k);
    Vector exp(CH);
    for (int i = 0; i < CH; ++i) {
      exp(i) = signs[k].expected(i);
    }
    expected.push_back(exp);
    Matrix cov(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
	cov(i, j) = signs[k].covariance(i, j);
      }
      cov(i, i) += noise;
    }
    lndet.push_back(log(fabs(cov.determinant())));
    inverse.push_back(cov.inverse());
  }

  SDataView<int> ren(size);
  vector<Cluster> ren_signs;
  ReadSegments(src, seg, ren, ren_signs);

  std::vector<int> ren_clu(ren_signs.size());
  std::vector<int> ren_ind(ren_signs.size());

  for (int s = 0; s < (int)ren_signs.size(); ++s) {
    std::cerr << s << " / " << ren_signs.size() << std::endl;
    Vector seg_exp(CH);
    for (int i = 0; i < CH; ++i) {
      seg_exp(i) = ren_signs[s].expected(i);
    }
    Matrix seg_mul(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
        seg_mul(i, j) = ren_signs[s].covariance(i, j) 
          + seg_exp(i) * seg_exp(j);
      }
    }
    
    int minPos = -1; 
    double minVal = 0.0;
    for (int k = 0; k < (int)index.size(); ++k) {
      double dist = mulTrace(inverse[k], seg_mul)
        - 2 * expected[k] * inverse[k] * seg_exp 
        + expected[k] * inverse[k] * expected[k];
      if (dist > leaveChiSquare) continue;
      dist += lndet[k];
      if (minPos == -1 || minVal > dist) {
        minPos = k;
        minVal = dist;
      }
    }
    ren_clu[s] = minPos != -1 ? index[minPos] : 0;
    ren_ind[s] = minPos;
  }

  int cr[256][256], r[256], s[256], fs[256];

  int pixcnt = 0, refcnt = 0;
  
  for (int i = 0; i < 256; ++i) {
    for (int j = 0; j < 256; ++j) {
      cr[i][j] = 0;
    }
    r[i] = 0; s[i] = 0; fs[i] = 0;
  }
  
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ren[i][j] == -1 || ren_clu[ren[i][j]] == 0) continue;
      ++fs[ren_clu[ren[i][j]]];
      ++pixcnt;
      if (ref[i][j] == 0) continue; 
      ++cr[ren_clu[ren[i][j]]][ref[i][j]];
      ++r[ref[i][j]];
      ++s[ren_clu[ren[i][j]]];
      ++refcnt;
    }
  }

  double refrate = (double)refcnt / (double)pixcnt;

  Todo todo[256];
  unsigned char best_index[256];

  std::vector<unsigned char> split_index[256];
  vector<Cluster> split_signs[256];

  int lea = 0, rec1 = 0, rec2 = 0;

  todo[0] = LEAVE;
  for (int i = 1; i < 256; ++i) {
    if (fs[i] == 0) {
      todo[i] = LEAVE;
    } else if (fs[i] * refrate >= s[i] * 8) {
      if (fs[i] > pixcnt * 0.005) {
        todo[i] = RECLASSIFY; ++rec1;
      } else {
        todo[i] = LEAVE; ++lea;
      }
    } else {
      int max = 0;
      for (int j = 0; j < 256; ++j) {
        if (max < cr[i][j]) {
          max = cr[i][j];
        }
      }
      if (max < s[i] * min_relevant) {
        todo[i] = RECLASSIFY; ++rec2;
      } else {
        vector<unsigned char> good;
        for (int j = 0; j < 256; ++j) {
          if (cr[i][j] >= max * min_rate) {
            good.push_back(j);
          }
        }
        if (good.size() == 1) {
          todo[i] = CLASSIFY;
          best_index[i] = good.front();
        } else {
          todo[i] = SPLIT;
          split_index[i] = good;
          split_signs[i].resize(good.size());
        }
      }
    }
  }

  std::cout << "Leave: " << std::count(todo + 0, todo + 256, LEAVE) << ' '
            << '(' << lea << ')' << std::endl;
  std::cout << "Classify: " << std::count(todo + 0, todo + 256, CLASSIFY) 
            << std::endl;
  std::cout << "Reclassify: " << std::count(todo + 0, todo + 256, RECLASSIFY) 
            << ' ' << '(' << rec1 << ',' << rec2 << ')' << std::endl;
  std::cout << "Split: " << std::count(todo + 0, todo + 256, SPLIT) 
            << std::endl;


  for (int s = 0; s < (int)ren_signs.size(); ++s) {
    unsigned char cluster = ren_clu[s];
    if (todo[cluster] == RECLASSIFY) {
      Vector seg_exp(CH);
      for (int i = 0; i < CH; ++i) {
        seg_exp(i) = ren_signs[s].expected(i);
      }
      Matrix seg_mul(CH, CH);
      for (int i = 0; i < CH; ++i) {
        for (int j = 0; j < CH; ++j) {
          seg_mul(i, j) = ren_signs[s].covariance(i, j) 
            + seg_exp(i) * seg_exp(j);
        }
      }
    
      int minPos = -1; 
      double minVal = 0.0;
      for (int k = 0; k < (int)index.size(); ++k) {
        if (todo[index[k]] != CLASSIFY && todo[index[k]] != SPLIT)
          continue;
        double dist = mulTrace(inverse[k], seg_mul)
          - 2 * expected[k] * inverse[k] * seg_exp
          + expected[k] * inverse[k] * expected[k];
        if (dist > leaveChiSquare) continue;
        dist += lndet[k];
        if (minPos == -1 || minVal > dist) {
          minPos = k;
          minVal = dist;
        }
      }
      ren_clu[s] = minPos != -1 ? index[minPos] : 0;
      ren_ind[s] = minPos;
    }
  }

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ren[i][j] != -1 && todo[ren_clu[ren[i][j]]] == SPLIT) {
        for (int k = 0; k < (int)split_index[ren_clu[ren[i][j]]].size(); ++k) {
          if (split_index[ren_clu[ren[i][j]]][k] == ref[i][j]) {
            split_signs[ren_clu[ren[i][j]]][k] += src[i][j];
            break;
          }
        }
      }
    }
  }

  vector<Vector> split_expected[256];
  vector<Matrix> split_inverse[256];
  vector<double> split_lndet[256];

  for (int k = 1; k < 256; ++k) {
    for (int l = 0; l < (int)split_index[k].size(); ++l) {
      Vector exp(CH);
      for (int i = 0; i < CH; ++i) {
        exp(i) = split_signs[k][l].expected(i);
      }
      split_expected[k].push_back(exp);
      Matrix cov(CH, CH);
      for (int i = 0; i < CH; ++i) {
        for (int j = 0; j < CH; ++j) {
          cov(i, j) = split_signs[k][l].covariance(i, j);
        }
        cov(i, i) += noise;
      }
      split_lndet[k].push_back(log(fabs(cov.determinant())));
      split_inverse[k].push_back(cov.inverse());
    }
  }

  vector<unsigned char> split_sub(ren_signs.size());

  for (int s = 0; s < (int)ren_signs.size(); ++s) {
    unsigned char cluster = ren_clu[s];
    if (todo[cluster] == SPLIT) {
      Vector seg_exp(CH);
      for (int i = 0; i < CH; ++i) {
        seg_exp(i) = ren_signs[s].expected(i);
      }
      Matrix seg_mul(CH, CH);
      for (int i = 0; i < CH; ++i) {
        for (int j = 0; j < CH; ++j) {
          seg_mul(i, j) = ren_signs[s].covariance(i, j) 
            + seg_exp(i) * seg_exp(j);
        }
      }
    
      int minPos = -1; 
      double minVal = 0.0;
      for (int k = 0; k < (int)split_index[cluster].size(); ++k) {
        double dist = mulTrace(split_inverse[cluster][k], seg_mul)
          - 2 * split_expected[cluster][k] * split_inverse[cluster][k] 
          * seg_exp 
          + split_expected[cluster][k] * split_inverse[cluster][k] 
          * split_expected[cluster][k];
        dist += split_lndet[cluster][k];
        if (minPos == -1 || minVal > dist) {
          minPos = k;
          minVal = dist;
        }
      }
      split_sub[s] = minPos;
    }
  }

  SDataView<unsigned char> trg(size);
  SDataView<unsigned char> prc(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      bool pc = false;
      prc[i][j] = 0;
      if (ren[i][j] == -1) {
        trg[i][j] = 0;
        pc = true;
        prc[i][j] = 1;
      } else if (todo[ren_clu[ren[i][j]]] == LEAVE) {
        trg[i][j] = 0;
        pc = true;
        prc[i][j] = 1;
      } else if (todo[ren_clu[ren[i][j]]] == CLASSIFY) {
        trg[i][j] = best_index[ren_clu[ren[i][j]]];
        
        Vector current(CH);
        for (int l = 0; l < CH; ++l) {
          current(l) = ((Point&)src[i][j])[l];
        }

        int cluster = ren_ind[ren[i][j]];

        Vector diff = current - expected[cluster];
        double dist = diff * inverse[cluster] * diff;
        if (dist > pixelLeaveChiSquare) {
          pc = true;
          prc[i][j] = 2;
        }
      } else if (todo[ren_clu[ren[i][j]]] == SPLIT) {
        trg[i][j] = split_index[ren_clu[ren[i][j]]][split_sub[ren[i][j]]];

        Vector current(CH);
        for (int l = 0; l < CH; ++l) {
          current(l) = ((Point&)src[i][j])[l];
        }

        int cluster = ren_ind[ren[i][j]];

        Vector diff = current - expected[cluster];
        double dist = diff * inverse[cluster] * diff;
        if (dist > pixelLeaveChiSquare) {
          pc = true;
          prc[i][j] = 3;
        }
        
        cluster = ren_clu[ren[i][j]];
        int sub = split_sub[ren[i][j]];
        diff = current - split_expected[cluster][sub];
        dist = diff * split_inverse[cluster][sub] * diff;
        if (dist > pixelLeaveChiSquare * 4.0) {
          pc = true;
          prc[i][j] = prc[i][j] == 3 ? 4 : 5;
        }        
      }
      if (pc) {
        if (((Point&)src[i][j]).zero()) continue;
        Vector current(CH);
        for (int l = 0; l < CH; ++l) {
          current(l) = ((Point&)src[i][j])[l];
        }
        int minPos = -1; 
        double minVal = 0.0;
        for (int k = 0; k < (int)index.size(); ++k) {
          if (todo[index[k]] != CLASSIFY && todo[index[k]] != SPLIT)
            continue;
          Vector diff = current - expected[k];
          double dist = diff * inverse[k] * diff;
          if (dist > leaveChiSquare) continue;
          dist += lndet[k];
          if (minPos == -1 || minVal > dist) {
            minPos = k;
            minVal = dist;
          }
        }
        int cluster = minPos != -1 ? index[minPos] : 0;

        if (todo[cluster] == LEAVE) {
          trg[i][j] = 0;
        } else if (todo[cluster] == CLASSIFY) {
          trg[i][j] = best_index[cluster];
        } else if (todo[cluster] == SPLIT) {
          Vector current(CH);
          for (int l = 0; l < CH; ++l) {
            current(l) = ((Point&)src[i][j])[l];
          }
          int minPos = -1; 
          double minVal = 0.0;
          for (int k = 0; k < (int)split_index[cluster].size(); ++k) {
            Vector diff = current - split_expected[cluster][k];
            double dist = diff * split_inverse[cluster][k] * diff;
            dist += split_lndet[cluster][k];
            if (minPos == -1 || minVal > dist) {
              minPos = k;
              minVal = dist;
            }
          }
          trg[i][j] = split_index[cluster][minPos];
        }
      }
    }
  }

  SFile trgfile(argv[5], SFile::CREATE);
  trgfile.createDataset("clasified", UINT8, size).
    setDataView<unsigned char>(trg);
  SFile prcfile("reclassify.hdf", SFile::CREATE);
  prcfile.createDataset("pixel reclassify", UINT8, size).
    setDataView<unsigned char>(prc);
  return 0;
}
